{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\16673\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\gym\\envs\\registration.py:555: UserWarning: \u001b[33mWARN: The environment SuperMarioBros-1-1-v0 is out of date. You should consider upgrading to version `v3`.\u001b[0m\n",
      "  logger.warn(\n",
      "c:\\Users\\16673\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\gym\\envs\\registration.py:627: UserWarning: \u001b[33mWARN: The environment creator metadata doesn't include `render_modes`, contains: ['render.modes', 'video.frames_per_second']\u001b[0m\n",
      "  logger.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[['NOOP'], ['right'], ['right', 'A'], ['right', 'B'], ['right', 'A', 'B'], ['A'], ['left']]\n",
      "[['NOOP'], ['right'], ['right', 'A'], ['right', 'B'], ['right', 'A', 'B'], ['A'], ['left'], ['left', 'A'], ['left', 'B'], ['left', 'A', 'B'], ['down'], ['up']]\n"
     ]
    }
   ],
   "source": [
    "import gym\n",
    "import time\n",
    "import random\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from collections import deque\n",
    "from gym.wrappers import GrayScaleObservation\n",
    "from nes_py.wrappers import JoypadSpace\n",
    "from gym_super_mario_bros.actions import SIMPLE_MOVEMENT\n",
    "from gym_super_mario_bros.actions import COMPLEX_MOVEMENT\n",
    "\n",
    "# 1. 创建马里奥环境\n",
    "env = gym.make(\"SuperMarioBros-1-1-v0\", apply_api_compatibility=True, render_mode=\"human\")\n",
    "env = JoypadSpace(env, SIMPLE_MOVEMENT)  # 约束动作空间\n",
    "env = GrayScaleObservation(env, keep_dim=True)  # 灰度化\n",
    "print(SIMPLE_MOVEMENT)\n",
    "print(COMPLEX_MOVEMENT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2. 定义 DQN 神经网络\n",
    "class MarioDQN(nn.Module):\n",
    "    def __init__(self, input_shape, num_actions):\n",
    "        super(MarioDQN, self).__init__()\n",
    "        self.conv = nn.Sequential(\n",
    "            nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(32, 64, kernel_size=4, stride=2),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(64, 64, kernel_size=3, stride=1),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(64 * 7 * 7, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, num_actions)\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.conv(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        return self.fc(x)\n",
    "\n",
    "# 3. 初始化 DQN 模型\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = MarioDQN((1, 84, 84), len(SIMPLE_MOVEMENT)).to(device)\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-4)\n",
    "loss_fn = nn.MSELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4. 经验回放（Replay Buffer）\n",
    "class ReplayBuffer:\n",
    "    def __init__(self, capacity):\n",
    "        self.buffer = deque(maxlen=capacity)\n",
    "\n",
    "    def push(self, state, action, reward, next_state, done):\n",
    "        self.buffer.append((state, action, reward, next_state, done))\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        return random.sample(self.buffer, batch_size)\n",
    "\n",
    "buffer = ReplayBuffer(10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5. 训练循环\n",
    "epsilon = 1.0  # 探索率\n",
    "gamma = 0.99  # 折扣因子\n",
    "\n",
    "for episode in range(1000):  # 训练 1000 局\n",
    "    state, _ = env.reset()\n",
    "    state = np.array(state, dtype=np.float32) / 255.0  # 归一化\n",
    "    done = False\n",
    "    total_reward = 0\n",
    "\n",
    "    while not done:\n",
    "        env.render()\n",
    "        time.sleep(0.01)  # 降低 CPU 负载\n",
    "        \n",
    "        # ε-greedy 选择动作\n",
    "        if random.random() < epsilon:\n",
    "            action = env.action_space.sample()  # 随机探索\n",
    "        else:\n",
    "            state_tensor = torch.tensor(state, device=device).unsqueeze(0)\n",
    "            with torch.no_grad():\n",
    "                action = model(state_tensor).argmax().item()\n",
    "        \n",
    "        # 执行动作\n",
    "        next_state, reward, done, truncated, info = env.step(action)\n",
    "        next_state = np.array(next_state, dtype=np.float32) / 255.0  # 归一化\n",
    "\n",
    "        # 奖励获取设定\n",
    "        # if action == 2:\n",
    "        #     reward += 0.7\n",
    "        # if action == 6 or action == 7:\n",
    "        #     reward -= 0.1\n",
    "\n",
    "        buffer.push(state, action, reward, next_state, done)\n",
    "        state = next_state\n",
    "        total_reward += reward\n",
    "\n",
    "    # 训练 DQN\n",
    "    if len(buffer.buffer) > 1000:\n",
    "        batch = buffer.sample(32)\n",
    "        state_batch = torch.tensor([b[0] for b in batch], dtype=torch.float32, device=device)\n",
    "        action_batch = torch.tensor([b[1] for b in batch], device=device)\n",
    "        reward_batch = torch.tensor([b[2] for b in batch], dtype=torch.float32, device=device)\n",
    "        next_state_batch = torch.tensor([b[3] for b in batch], dtype=torch.float32, device=device)\n",
    "        done_batch = torch.tensor([b[4] for b in batch], dtype=torch.float32, device=device)\n",
    "\n",
    "        q_values = model(state_batch).gather(1, action_batch.unsqueeze(1)).squeeze()\n",
    "        next_q_values = model(next_state_batch).max(1)[0]\n",
    "        expected_q_values = reward_batch + (1 - done_batch) * gamma * next_q_values\n",
    "\n",
    "        loss = loss_fn(q_values, expected_q_values.detach())\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    # 逐步降低探索率\n",
    "    epsilon = max(0.1, epsilon * 0.995)\n",
    "\n",
    "    print(f\"Episode {episode}, Total Reward: {total_reward}\")\n",
    "\n",
    "    # 每 50 局保存模型\n",
    "    if episode % 50 == 0:\n",
    "        torch.save(model.state_dict(), \"mario_dqn.pth\")\n",
    "\n",
    "env.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
